from ..errors import BPlusTreeError


# TODO: 把这个node可以序列化为slotted page的字节集bytes
# TODO: LSN的实现 -> 被修改的节点，需要显性标记一下LSN号
class BPlusTreeNode:
    def __init__(self, is_leaf=True):
        # 是否是叶子节点, true表示是叶子节点
        self.is_leaf = is_leaf
        self.keys = []
        self.children = []
        # 只有在叶子节点才有value
        self.values = []

        # 如果我们通过指针进行回溯，那么我们应该有下述字段
        # self.parent = 
        self.next_leaf = None
        self.lsn = 0

    def get_child(self, i):
        # key ->         [1, 3, 5]
        # children ->  [>1,1-3, 3-5, >=5]
        if i >= len(self.children):
            return self.next_leaf
        return self.children[i]

    def count_childern(self):
        if self.next_leaf:
            return len(self.children) + 1
        return len(self.children)



class BPlusTree:
    def __init__(self):
        self.root = BPlusTreeNode(is_leaf=True)
        # 暂时使用disk变量表示磁盘，后面会具体来实现这部分
        self.disk = {}


    def insert(self, key, value):
        # 实现B+树的插入操作
        if key is None:
            raise BPlusTreeError("invalid key")
        
        # 直接插入叶子节点中
        node = self.find_leaf_node(key)
        # 正因为，调用了下面的函数，我们可以保证，插入过程是
        # 有序的，因为该函数，寻找的是最右边的相同的key的下标，
        # 如果没有找到 0
        index = self._find_rightmost_key_index(node, key)
        node.keys.insert(index, key)
        node.values.insert(index, value)

        # 分裂，也就是不断递归，向父节点插入元素的过程
        if self._need_split(node):
            self._split(node)

    def delete(self, key, value=None):
        node = self.find_leaf_node(key)
        while node:
            indexes = list(self._find_indexes(node.keys, key))
            if len(indexes) == 0:
                break
            # 如何去解决删除一个元素后，对应的下标移位的问题
            i = 0
            deleted_count = 0
            while i < len(indexes):
                index = indexes[i]
                actual_index = index - deleted_count
                # 跳过 value 不等于参数的 key
                if value is not None:
                    if node.values[actual_index] != value:
                        i += 1
                        continue
                node.keys.pop(actual_index)
                node.values.pop(actual_index)
                i += 1
                deleted_count += 1

            node = node.next_leaf
        # 准确来说，此时还应该补充一个合并机制(coalesce)，即把小于 n/2 的node进行合并

    def find_range(self, start=float('-inf'), end=float('inf'), return_keys=False):
        # select * from t1 where a > 100;
        # 不包含等值
        # 因为，对于等值，可以直接补充等值查询即可
        # 算法复杂度没有新增多少
        values = []
        node = self.find_leaf_node(start)
        while node:
            for index, key in enumerate(node.keys):
                # 提前退出
                if key >= end:
                    break
                # 如果我们不在上面指定trick start=-inf, ...
                # 那么我们就要判断 start/end 是否为 None
                if start < key < end:
                    if return_keys:
                        values.append(key)
                    else:
                        values.append(node.values[index])
            node = node.next_leaf
        return values

    def _split(self, node: BPlusTreeNode):
        """ 用于调整B+树的结构，用于做节点的分裂 """
        middle_index = len(node.keys) // 2
        # 把当前的 node 节点，拆分成相等元素的两个节点
        # 这种在工程上有不同的实现(改良)
        right_node = BPlusTreeNode(is_leaf=node.is_leaf)
        left_node = node

        # 拆分完之后，就要导数据
        # 新节点就是右节点，原来的旧节点就是左节点
        # 我们这里面之所以复用原来的节点，是因为传入的参数是一个引用(指针)
        # 如果直接用新的节点进行替换，就会出现找不到节点的问题
        right_node.keys.extend(node.keys[middle_index:])
        right_node.children.extend(node.children[middle_index:])
        right_node.values.extend(node.values[middle_index:])
        right_node.next_leaf = node.next_leaf

        left_node.keys = node.keys[:middle_index]
        left_node.children = node.children[:middle_index]
        left_node.values = node.values[:middle_index]
        left_node.next_leaf = right_node

        assert (len(left_node.keys) > 0) and (len(right_node.keys) > 0)
        assert tuple(left_node.keys) < tuple(right_node.keys)

        if node is self.root:
            new_root = BPlusTreeNode(is_leaf=False) # 肯定不是一个叶子节点
            new_root.keys.append(right_node.keys[0])
            new_root.children.extend([left_node, right_node])
            self.root = new_root
        else:
            parent = self._find_parent(self.root, node)
            index = parent.children.index(node)
            parent.keys.insert(index, right_node.keys[0])
            # parent.children[index] = left_node
            parent.children.insert(index + 1, right_node)

            if self._need_split(parent):
                self._split(parent)

    def find_leaf_node(self, key):
        """ 寻找最左边的叶子节点(我们B+树是按照从小到大组织数据结构的) """
        node = self.root
        while node and not node.is_leaf:
            index = self._find_leftmost_key_index(node, key)
            # 如果没有这样的 index，则 node 为最后一个节点
            if index >= len(node.keys):
                node = node.children[-1]

            # 对于唯一的key，应该使用下描述方法，但是不唯一的key
            # 会很麻烦，不能直接使用下述代码
            # elif key == node.keys[index]:
            #     node = node.get_child(index+1)
            else:
                node = node.get_child(index)

        while node.keys and node.keys[-1] < key and node.next_leaf:
            node = node.next_leaf
        return node

        
    @staticmethod
    def _need_split(node):
        return len(node.keys) > 10 # 这里是随便写的，意思是B+树节点中的元素数量如果大于这个值就分裂


    @staticmethod
    def _find_leftmost_key_index(node: BPlusTreeNode, key):
        # 寻找 key <= node.keys 中最小值时的下标
        # 其他更高效的方法：二分查找，但二分查找不能解决元素相同的问题
        i = 0
        while i < len(node.keys):
            if key <= node.keys[i]:
                return i
            i += 1
        return len(node.keys)

    @staticmethod
    def _find_rightmost_key_index(node: BPlusTreeNode, key):
        # 实现查找插入位置
        # 假如：
        # node.keys：[1, 3, 10,          100]
        #           /  |  |    |            \
        #       [-1,0][2][4] [11,15,99]    [101] -> children
        i = 0
        # 遍历node.keys，找到第一个大于key的索引
        while i < len(node.keys):
            if node.keys[i] <= key:
                i += 1
            else:
                return i
        # 如果没有找到，返回node.keys的长度
        return len(node.keys)

        
    # 实现查找父节点的操作
    def _find_parent(self, current_node: BPlusTreeNode, target_node: BPlusTreeNode) -> BPlusTreeNode:
        # 如果当前节点是叶子节点或者目标节点在当前节点的子节点中，则返回当前节点
        if current_node.is_leaf or target_node in current_node.children:
            return current_node

        # 遍历当前节点的子节点
        for child in current_node.children:
            # 如果目标节点在子节点的子节点中，则递归调用_find_parent函数查找父节点
            if target_node in child.children:
                return self._find_parent(child, target_node)
            
        # 如果没有找到父节点，则返回None
        return None

    def find(self, key):
        values = []
        node = self.find_leaf_node(key)
        while node:
            indexes = list(self._find_indexes(node.keys, key))
            if len(indexes) == 0:
                break
            for index in indexes:
                values.append(node.values[index])
            node = node.next_leaf
        return values
    
    @staticmethod
    def _find_indexes(keys, key):
        # 遍历keys列表
        for i, k in enumerate(keys):
            # 如果key等于k，则返回i
            if key == k:
                yield i
            # 这样可以做个减枝叶
            elif key < k:
                break
